#/*************************************************************************
#    > File Name: Lib/ARCH_Operation.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: 2022年05月26日 星期四 15时22分37秒
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from File_Operation import *
from Folder_Operation import *
from BASDA_Operation import *
from List_Operation import *
from Value_Calculation import *

# 生成ARCH的运行参数

# 按mg的比例计算 bps 给出的加了parton level cut 的截面大小
def Cal_Parton_Cut_Xection(cut_file_name, nocut_file_name, theoretical_parafile_name):
	#生成MG对应的加parton_cut
	cut_mg5 = pd.read_table(cut_file_name, delim_whitespace=True)    #读取.txt文件
	if len(cut_mg5.columns)==1:
		cut_mg5 = pd.read_table(cut_file_name, sep=",")    #读取.txt文件

	cut_xection=cut_mg5["cross"]

	nocut_mg5 = pd.read_table(nocut_file_name, delim_whitespace=True)
	if len(nocut_mg5.columns)==1:
		nocut_mg5 = pd.read_table(nocut_file_name, sep=",")
	no_cut_xection=nocut_mg5["cross"]

	df_theo = pd.read_table(theoretical_parafile_name)
	no_cut_theoretical=df_theo["xsh_W4b"]*0.2134

	cut_cross=[round(xx,2) for xx in np.multiply(cut_xection,no_cut_theoretical)/no_cut_xection]   #np.multiply()数组和矩阵对应位置相乘，输出与相乘数组/矩阵的大小一致
	return(cut_cross)


# 更改 ARCH contraol文件中参数
def Replace_Para_in_ARCH(input_para, cut_cross):
	df_num=len(cut_cross)

	output_list={}
	for key in input_para:
		param_list_all=[]
		for cross_num, cross_data in enumerate(cut_cross):
			values={}
			format_num=str(cross_num+1).rjust(2,"0")
			values={"fnum":format_num,"fvalue":str(cross_data)}
			input_dict=Trans_MultiValue_to_Dict(input_para[key],values)   #替换,values传参到input_para[key]中
			input_list=Trans_Dict_to_List(input_dict)   #将dict转换为list
			param_list_all.append(input_list)
		output_list[key]=param_list_all

	return(output_list)



# 生成ARCH的control文件， 
# 其中cut_cross 由Cal_Parton_Cut_Xection计算
# dir_name 是一个列表， 第一个为control文件模板，第二个为复制后使用的文件
# para_path 为修改的path的改动参数
def Generate_Multi_ARCH_Control_File(cut_cross,dir_name, input_para):
	
	target_folder=Copy_File(dir_name["source"],dir_name["target"])   #复制文件
	print("In Generate_Multi_ARCH_Control_File, the generated ARCH control files are in ", target_folder)
# generate param files

	input_origin={}
	for key in input_para:
		input_origin[key] = input_para[key]["origin"]
	replaced_para=Replace_Para_in_ARCH(input_origin, cut_cross)   #替换input_origin中的cut_cross
	for key in input_para:
		File_Replace_with_Para(replaced_para[key], input_para[key]["source"], target_folder, input_para[key]["labels"])    #生成多个替换文件,替换input_para[key]["labels"]到replaced_para[key]中

	return(replaced_para)



# 利用ARCH的结果



# 从 ARCH 输出的tmp文件中，提取截面相关信息
# 返回值为
def Find_ARCH_xection(ARCH_out_name,total_file_num, ARCH_summary_folder):
	target_folder = Create_Folder_from_FileDir(ARCH_summary_folder,"/")   #创建文件夹,去掉基文件
	print("target_folder",target_folder)
	ARCH_xection = []
	ARCH_number= []
	pos=[]
	for file_num in range(total_file_num):
		format_num=str(file_num+1).rjust(2,"0")
		ARCH_file = ARCH_out_name.format(format_num)
		#label = {'OutPut --- Efficience for Cut': 17}
		label = {'total cross section is': 15}
		word_list, input_line_list = Find_Line(ARCH_file, label)   #查找ARCH_file中的label
		strings = ''.join(word_list)
		out_file = ARCH_summary_folder + '/ARCH_output_{}.dat'.format(format_num)
		#print(out_file)
		f = open(out_file, 'w')
		f.write(strings)
		f.close()
		with open(out_file, 'r') as (file):
			yfile = yaml.safe_load(file)
			strings = yfile['number passed cut basic_lep is']
			key = 'xection'
			if strings.count(key) >= 1:
				xection = strings.split(',',3)[2].split(' ',4)[3]
				number = strings.split(',',3)[0]
				ARCH_xection.append(xection)
				ARCH_number.append(number)
		pos.append(file_num)
	cut_cross=[float(value) for value in ARCH_xection]
	cut_number=[int(value) for value in ARCH_number]
	pos_number=[int(value) for value in pos]
	ARCH_results= {"cross":cut_cross,"number":cut_number,"pos":pos_number}
	return(ARCH_results)


# 将 ARCH 中得到的截面提取出来，放入列表  cut_cross 中
def Get_Cross_after_ARCH(ARCH_tmp_name, total_file_num,ARCH_summary_folder):
	target_folder = Create_Folder(ARCH_summary_folder,"/")
	#print("in Get_Cross_after_ARCH: target_folder is",target_folder)
	ARCH_results=Find_ARCH_xection(ARCH_tmp_name,total_file_num, ARCH_summary_folder)
	return(ARCH_results)


def Collect_ARCH_Cross(theoretical_para_file_name,para_name, ARCH_results,ARCH_param_out_file):
	target_folder = Create_Folder_from_FileDir(ARCH_param_out_file,"/")
	origin_df=Load_THDM_Param(theoretical_para_file_name,para_name)
	theo_df=origin_df.loc[ARCH_results["pos"]]
	ARCH_df=pd.DataFrame(ARCH_results)
	aRCH_df=ARCH_df.set_index("pos")
	ARCH_df=pd.concat([theo_df,ARCH_df],axis=1)
	ARCH_df.to_csv(ARCH_param_out_file)
	return(ARCH_results)
